import gym
import numpy as np
import pandas as pd

SIZE = 4
ACTION_PROB = 0.25
GAMMA = 0.99


class MyGridEnv(gym.Env):
    metadata = {
        'render.modes': ['human', 'rgb_array'],
        'video.frames_per_second': 2
    }

    def __init__(self):
        self.nA = 4
        self.nS = 16
        self.P = {}
        self.gamma = GAMMA
        self.__state = np.random.randint(1, SIZE * SIZE + 1)
        self.dict = dict()
        self.observation_space = range(1, SIZE * SIZE + 1)
        self.__terminal_space = dict()

        self.__terminal_space[1] = 1
        self.__terminal_space[16] = 1

        self.action_space = ('n', 'e', 's', 'w')
        self._action = {'n': np.array([0, 1]),
                        'e': np.array([1, 0]),
                        's': np.array([0, -1]),
                        'w': np.array([-1, 0])}

        tmp = 1
        for i in range(SIZE):
            for j in range(SIZE):
                self.dict[(i, j)] = tmp
                self.dict[tmp] = (i, j)
                tmp += 1

        self.viewer = None

        def get_next_state_and_reward(s, a):
            if s in self.__terminal_space:
                return s, self._reward(s)

            cur_x, cur_y = self.dict[s]
            new_pos = tuple((np.array((cur_x, cur_y)) + self._action[a]).tolist())
            new_x, new_y = new_pos
            new_state = None
            if new_x < 0 or new_x >= SIZE or new_y < 0 or new_y >= SIZE:
                new_state = s
            else:
                new_state = self.dict[new_pos]

            r = self._reward(new_state)
            is_terminal = False
            if new_state in self.__terminal_space:
                is_terminal = True
            return new_state, r


        self.state_pos = {}
        for i in range(1, SIZE * SIZE + 1):
            y = (500 - ((i - 1) // 4) * 100 - 50)
            left = (i - 1) % 4
            x = 100 + 50 + left * 100
            self.state_pos[i] = (x, y)


        for s in range(1, SIZE * SIZE + 1):
            # envs.P[s][a] is a list of transition tuples (prob, next_state, reward, done).
            self.P[s] = {}
            for a in self.action_space:
                next_s, r = get_next_state_and_reward(s, a)
                self.P[s][a] = (1, next_s, r, next_s in self.__terminal_space)

    def reset(self):
        self.__state = np.random.randint(1, SIZE * SIZE + 1)
        return self.__state

    def _reward(self, state):
        r = -1.0
        if state in (1, 16):
            r = 0   # 这里终止状态的奖励不能是1，不然在P里面很不好弄，而且迭代无法收敛
        return r

    def transform(self, state, action):
        if state in self.__terminal_space:
            return state, self._reward(state), True, {}

        cur_x, cur_y = self.dict[state]
        new_pos = tuple((np.array((cur_x, cur_y)) + action).tolist())
        new_x, new_y = new_pos
        new_state = None
        if new_x < 0 or new_x >= SIZE or new_y < 0 or new_y >= SIZE:
            new_state = state
        else:
            new_state = self.dict[new_pos]

        r = self._reward(new_state)
        is_terminal = False
        if new_state in self.__terminal_space:
            is_terminal = True
        return new_state, r, is_terminal, {}

    def step(self, action_str):
        state = self.__state
        action_ = self._action[action_str]
        next_state, r, is_terminal, _ = self.transform(state, action_)
        self.__state = next_state
        return next_state, r, is_terminal, {}

    def render(self, mode='human', close=False):
        if close:
            if self.viewer is not None:
                self.viewer.close()
                self.viewer = None
            return
        from gym.envs.classic_control import rendering
        screen_width = 600
        screen_height = 600

        if self.viewer is None:
            self.viewer = rendering.Viewer(screen_width, screen_height)

            # 创建网格世界
            self.line1 = rendering.Line((100, 100), (500, 100))
            self.line2 = rendering.Line((100, 200), (500, 200))
            self.line3 = rendering.Line((100, 300), (500, 300))
            self.line4 = rendering.Line((100, 400), (500, 400))
            self.line5 = rendering.Line((100, 500), (500, 500))
            self.line6 = rendering.Line((100, 100), (100, 500))
            self.line7 = rendering.Line((200, 100), (200, 500))
            self.line8 = rendering.Line((300, 100), (300, 500))
            self.line9 = rendering.Line((400, 100), (400, 500))
            self.line10 = rendering.Line((500, 100), (500, 500))

            self.start = rendering.make_polygon([(100, 400), (200, 400),
                                                 (200, 500), (100, 500)], filled=False)
            self.start.set_color(0.5, 0.6, 0.7)
            self.end = rendering.make_polygon([(400, 100), (500, 100),
                                                 (500, 200), (400, 200)], filled=False)
            self.end.set_color(0.5, 0.6, 0.7)

            self.agent = rendering.make_circle(40)
            self.circletrans = rendering.Transform(translation=self.state_pos[self.__state])
            self.agent.add_attr(self.circletrans)
            self.agent.set_color(0.1, 0.6, 0.1)


            self.line1.set_color(0, 0, 0)
            self.line2.set_color(0, 0, 0)
            self.line3.set_color(0, 0, 0)
            self.line4.set_color(0, 0, 0)
            self.line5.set_color(0, 0, 0)
            self.line6.set_color(0, 0, 0)
            self.line7.set_color(0, 0, 0)
            self.line8.set_color(0, 0, 0)
            self.line9.set_color(0, 0, 0)
            self.line10.set_color(0, 0, 0)

            self.viewer.add_geom(self.line1)
            self.viewer.add_geom(self.line2)
            self.viewer.add_geom(self.line3)
            self.viewer.add_geom(self.line4)
            self.viewer.add_geom(self.line5)
            self.viewer.add_geom(self.line6)
            self.viewer.add_geom(self.line7)
            self.viewer.add_geom(self.line8)
            self.viewer.add_geom(self.line9)
            self.viewer.add_geom(self.line10)
            self.viewer.add_geom(self.start)
            self.viewer.add_geom(self.end)
            self.viewer.add_geom(self.agent)

        if self.__state is None:
            return None

        self.circletrans.set_translation(self.state_pos[self.__state][0], self.state_pos[self.__state][1])
        return self.viewer.render(return_rgb_array=(mode == 'rgb_array'))

    def set_state(self, x):
        self.__state = x


if __name__ == '__main__':
    env = MyGridEnv()
    # envs.reset()
    # envs.render()

    env.set_state(15)
    env.render()

    pass
